"""
@Author: CheyenneLam
@Date: 2020-12-16 10:31:00
@LastEditTime: 2020-12-16 10:31:00
@LastEditors: CheyenneLam
@Description: 封装数据库连接方法(mysql,pg,mongo,hana)
@FilePath: \Main\sql.py
"""
import os
from odps import ODPS
from odps.df import DataFrame
import psycopg2
import configparser
import pandas as pd
from pymongo import MongoClient
import pymysql

# dataworks和数据库映射关系
projectMapping = {
    '1': ['data_statistic', 'pg_op', 'mongo_op', 'dDing_data_statistic'],
    '2': ['work_prod', 'pg_op', 'mongo_op', 'dDing_prod'],
    '3': ['work_pre', 'pg_pre', 'mongo_pre', 'dDing_pre'],
    '4': ['work_dev', 'pg_test', 'mongo_test', 'dDing_test'],
}


def get_config(section):
    """
    获取连接数据库的配置
    """
    config = configparser.ConfigParser()
    path = os.path.join(os.path.dirname(__file__), 'conf/db.conf')
    config.read(path)
    return dict(config.items(section))


class OdpsDB():
    """
    dataworks有关连接、查询
    """

    def __init__(self, project_name):
        self.projectName = project_name
        self.configDict = get_config('dataworks')
        self.o = ODPS(self.configDict['ak'], self.configDict['sk'], self.projectName,
                      endpoint=self.configDict['endpoint'])

    def query_sql(self, strsql, project_name=None):
        """
        从dataworks读取数据
        """
        if project_name is None:
            project_name = self.projectName
        instance = self.o.execute_sql(strsql, project_name)
        with instance.open_reader() as reader:
            res = reader.to_pandas()
        return res

    def count_table(self, table_name, project_name=None):
        """
        从dataworks读取数据条数
        """
        if project_name is None:
            project_name = self.projectName
        res = DataFrame(self.o.get_table(table_name, project=project_name)).count().execute()
        return res


class PgDB():
    """
    PostgreSQL有关连接、查询
    """

    def __init__(self, pg_config):
        self.configDict = pg_config
        self.server = None
        self.localPort = int(self.configDict['port'])
        # 进行连接
        self.conn = psycopg2.connect(database=self.configDict['db_name'],
                                     user=self.configDict['username'],
                                     password=self.configDict['password'],
                                     host=self.configDict['uri'],  # host、port 固定
                                     port=self.localPort)

    def query_sql(self, strsql):
        res = pd.read_sql(strsql, self.conn)
        return res

    def close(self):
        """关闭连接"""
        self.conn.close()


class MongoDB():
    """
     mongoDB 连接，查询
    """
    def __init__(self, mongo_config):
        self.configDict = mongo_config
        self.server = None
        self.localPort = int(self.configDict['port'])

        self.client = MongoClient(
            host=self.configDict['uri'],
            port=self.localPort
        )

        if self.configDict["username"] == "root":
            self.database = self.client.admin
            self.database.authenticate(self.configDict['username'], self.configDict['password'])
            self.database = self.client[self.configDict['db_name']]
        else:
            self.database = self.client[self.configDict['db_name']]
            self.database.authenticate(self.configDict['username'], self.configDict['password'])

    def query_collection_count(self, collection_name):
        collection_count = self.database.get_collection(collection_name).find().count()
        return collection_count

    def close(self):
        self.client.close()

class MysqlDB():
    """
    连接mysql数据库以及读取数据
    """
    def __init__(self, dbName):
        """
        :param dbName: 数据库名
        :description: 初始化数据库连接信息
        """
        self.configDict = get_config(dbName)
        self.ip = self.configDict['uri']
        self.port = int(self.configDict['port'])
        self.username = self.configDict['username']
        self.password = self.configDict['password']
        self.database = self.configDict['db_name']
        self.conn = pymysql.connect(user = self.username, passwd = self.password,host = self.ip, database=self.database, port=self.port)
    
    def query_sql(self, sqlStr):
        """
        :param sqlStr: 查询的sql语句
        :description: 查询数据
        """
        res = pd.read_sql(sqlStr, con = self.conn)
        return res
    
    def close(self):
        """
        :description: 关闭连接
        """
        self.conn.close()

class SapHanaDB():
    """
    连接Hana数据库以及读取数据
    """
    def __init__(self, dbName):
        """
        :param dbName: 数据库名
        :description: 初始化数据库连接信息
        """
        self.configDict = get_config(dbName)
        self.ip = self.configDict['uri']
        self.port = int(self.configDict['port'])
        self.username = self.configDict['username']
        self.password = self.configDict['password']
        self.conn = pyhdb.connect(host=self.ip, port=self.port, user=self.username, password=self.password)
    
    def query_sql(self, sqlStr):
        """
        :param sqlStr: 查询的sql语句
        :description: 查询数据
        """
        res = pd.read_sql(sqlStr, con = self.conn)
        return res
    
    def close(self):
        """
        :description: 关闭连接
        """
        self.conn.close()